import os
import json
import typer
from typing import Optional
from datasets import load_dataset
from intelli_port.commons.config import load_settings
from .registry import ensure_dir, save_metadata


app = typer.Typer()


@app.command()
def run(
    data_dir: Optional[str] = None,
    output_dir: Optional[str] = None,
    lora_r: int = 16,
    lora_alpha: int = 64,
    lora_dropout: float = 0.05,
    lr: float = 0.0001,
    seq_len: int = 2048,
    epochs: int = 3,
    grad_accum: int = 32,
    warmup_ratio: float = 0.05,
):
    s = load_settings()
    data_dir = data_dir or s.qwen_sft_data_dir
    output_dir = output_dir or s.qwen_output_dir
    ensure_dir(output_dir)
    files = [os.path.join(data_dir, x) for x in os.listdir(data_dir) if x.endswith(".jsonl")]
    count = 0
    for fp in files:
        ds = load_dataset("json", data_files=fp, split="train")
        count += len(ds)
    cfg = {
        "lora_r": lora_r,
        "lora_alpha": lora_alpha,
        "lora_dropout": lora_dropout,
        "lr": lr,
        "seq_len": seq_len,
        "epochs": epochs,
        "grad_accum": grad_accum,
        "warmup_ratio": warmup_ratio,
        "samples": count,
        "data_dir": data_dir,
    }
    with open(os.path.join(output_dir, "last_config.json"), "w", encoding="utf-8") as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)
    save_metadata(output_dir, "qwen", "qlora", output_dir, {"samples": count})
    typer.echo(output_dir)


if __name__ == "__main__":
    app()